package info.wangyuan.agent.service.rag;


import info.wangyuan.agent.entity.dto.RagDocSyncChunkMsgDTO;
import info.wangyuan.agent.entity.po.DocumentUnit;
import info.wangyuan.agent.mapper.DocumentUnitMapper;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.Map;

/**
 * @author Albert
 * @since 2025-08-30 21:45:06
 */
public interface RagDocSyncChunkStrategy {

    void handle(RagDocSyncChunkMsgDTO ragDocSyncChunkMsgDTO, String strategy) throws Exception;

    /**
     * 文件转字节
     */
    default byte[] fileToBytes(String filePath) throws IOException {
        File file = new File(filePath);
        try (FileInputStream fis = new FileInputStream(file); ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
            byte[] buffer = new byte[4096];
            int bytesRead;
            while ((bytesRead = fis.read(buffer)) != -1) {
                outputStream.write(buffer, 0, bytesRead);
            }
            return outputStream.toByteArray();
        }
    }

    /**
     * 保存数据
     *
     * @param ragDocSyncChunkMsgDTO 消息数据
     * @param ocrData               ocr数据
     */
    default void insertData(RagDocSyncChunkMsgDTO ragDocSyncChunkMsgDTO, Map<Integer, String> ocrData, DocumentUnitMapper documentUnitMapper) {

        for (int pageIndex = 0; pageIndex < ragDocSyncChunkMsgDTO.getPageSize(); pageIndex++) {

            String content = ocrData.getOrDefault(pageIndex, null);

            final DocumentUnit documentUnit = DocumentUnit.builder()
                    .fileId(ragDocSyncChunkMsgDTO.getFileId())
                    .pageIndex(pageIndex)
                    .content(content)
                    .isChunked(true)
                    .isVectored(false)
                    .build();

            if (content == null) {
                documentUnit.setIsChunked(false);
            }

            documentUnitMapper.insert(documentUnit);
        }
    }

}
